package com.canary;

import com.alibaba.cloud.nacos.ribbon.NacosServerIntrospector;
import com.netflix.loadbalancer.ILoadBalancer;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import org.apache.commons.lang.StringUtils;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.cloud.netflix.ribbon.ServerIntrospector;
import org.springframework.stereotype.Component;

import java.util.LinkedHashSet;
import java.util.List;
import java.util.stream.Collectors;

/**
 * 灰度规则
 *
 * @author dzh
 * @date 2023/1/4 8:26
 */
@Component
@ConditionalOnProperty(name = "canary.rule", havingValue = "canary")
public class CanaryRule implements IRule {
    public static final String CANARY_KEY = "canary";
    ILoadBalancer loadBalancer;
    LinkedHashSet<Server> serverList = new LinkedHashSet<>();
    int index = 0;
    @Value("${spring.cloud.nacos.discovery.group}")
    String canary;
    ServerIntrospector serverIntrospector = new NacosServerIntrospector();

    @Override
    public Server choose(Object key) {
        filterServer();
        if (serverList.size() == 0) {
            return null;
        }
        return serverList.stream().skip(index++ % serverList.size()).findFirst().orElse(null);
    }

    void filterServer() {

        if (serverList.size() > 1) {
            return;
        }

        boolean hasCanary = loadBalancer.getReachableServers().stream()
                .anyMatch(server -> StringUtils.equals(serverIntrospector.getMetadata(server).get(CANARY_KEY), canary));

        // 默认获取没有指定虚拟环境的服务
        String[] canaryCompare = new String[]{""};
        if (hasCanary) {
            // 如果有同名指定的虚拟环境，则选择
            canaryCompare[0] = canary;
        }
        List<Server> serverList = loadBalancer.getReachableServers().stream()
                .filter(server -> StringUtils.equals(serverIntrospector.getMetadata(server).get(CANARY_KEY), canaryCompare[0]))
                .collect(Collectors.toList());

        // 选择结果为同名灰度环境或者非灰度环境的服务
        this.serverList.addAll(serverList);
    }

    @Override
    public void setLoadBalancer(ILoadBalancer lb) {
        loadBalancer = lb;
    }

    @Override
    public ILoadBalancer getLoadBalancer() {
        return loadBalancer;
    }
}
